package cn.wpy.web;

import cn.wpy.service.AIService;
import com.alibaba.dashscope.common.Message;
import lombok.RequiredArgsConstructor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.ArrayList;
import java.util.List;

@RequiredArgsConstructor
public class MyWebSocketHandler extends TextWebSocketHandler {

    private static final Logger logger = LoggerFactory.getLogger(MyWebSocketHandler.class);
    private final AIService aiService;

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        logger.info("Connected to client");
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String userMessage = message.getPayload();
        List<Message> messages = new ArrayList<>();
        messages.add(Message.builder().role("system").content("你是一位有帮助的助手").build());
        messages.add(Message.builder().role("user").content(userMessage).build());

        aiService.streamCallWithCallback(messages, content -> {
            try {
                session.sendMessage(new TextMessage(content));
            } catch (Exception e) {
                logger.error("Error sending message", e);
            }
        });
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        logger.info("Connection closed");
    }
}